"""
Jupyter Notebook  - Proposal Threshold Optimization

 Jupyter Notebook 
"""

import os
import json
from pathlib import Path
from typing import Optional, Tuple, Dict
import matplotlib.pyplot as plt
import seaborn as sns
from PIL import Image
from IPython.display import display, Markdown, HTML


class NotebookThresholdOptimizer:
    """Jupyter Notebook """

    def __init__(self):
        """"""
        from src.config import DATA_YAML, STAGE1_CONFIG
        self.data_yaml = DATA_YAML
        self.stage1_weights = STAGE1_CONFIG['weights_path']
        self.results = None
        self.output_dir = 'proposal_threshold_optimization'

    def run_optimization(
        self,
        threshold_range: Tuple[float, float] = (0.001, 0.1),
        num_points: int = 15,
        sample_size: Optional[int] = None,
        target_min: int = 20,
        target_max: int = 100,
        output_dir: str = 'proposal_threshold_optimization',
        verbose: bool = True
    ) -> Dict:
        """
        

        Args:
            threshold_range: 
            num_points: 
            sample_size: 
            target_min: proposals
            target_max: proposals
            output_dir: 
            verbose: 

        Returns:
            
        """
        from src.optimization.optimize_proposal_threshold import run_optimization

        if verbose:
            print("  Proposal ")
            print("=" * 70)
            print(f": {threshold_range}")
            print(f": {num_points}")
            print(f": {sample_size or ''}")
            print(f" proposals: {target_min}-{target_max}")
            print("\n⏱ : 20-60 ")
            print("...\n")

        self.output_dir = output_dir

        # 
        results = run_optimization(
            data_yaml=self.data_yaml,
            stage1_weights=self.stage1_weights,
            threshold_range=threshold_range,
            num_points=num_points,
            sample_size=sample_size,
            output_dir=output_dir,
            target_min=target_min,
            target_max=target_max
        )

        self.results = results

        if verbose:
            print("\n" + "=" * 70)
            print(" !")
            print("=" * 70)

        return results

    def display_best_result(self):
        """"""
        if not self.results:
            print(" ")
            return

        best = self.results['best_metrics']

        # HTML
        html = f"""
        <div style="background-color: #f0f8ff; padding: 20px; border-radius: 10px; border: 2px solid #4CAF50;">
            <h2 style="color: #4CAF50;"> </h2>
            <table style="width: 100%; margin-top: 10px;">
                <tr style="background-color: #4CAF50; color: white;">
                    <th style="padding: 10px; text-align: left;"></th>
                    <th style="padding: 10px; text-align: right;"></th>
                </tr>
                <tr style="background-color: #ffffff;">
                    <td style="padding: 10px;"><strong></strong></td>
                    <td style="padding: 10px; text-align: right; font-size: 18px; color: #4CAF50;"><strong>{best['threshold']:.4f}</strong></td>
                </tr>
                <tr style="background-color: #f9f9f9;">
                    <td style="padding: 10px;">Recall @ IoU=0.5</td>
                    <td style="padding: 10px; text-align: right;">{best['recall_at_05']:.4f} ({best['recall_at_05']*100:.2f}%)</td>
                </tr>
                <tr style="background-color: #ffffff;">
                    <td style="padding: 10px;">Recall @ IoU=0.75</td>
                    <td style="padding: 10px; text-align: right;">{best['recall_at_075']:.4f} ({best['recall_at_075']*100:.2f}%)</td>
                </tr>
                <tr style="background-color: #f9f9f9;">
                    <td style="padding: 10px;">Mean Max IoU</td>
                    <td style="padding: 10px; text-align: right;">{best['mean_max_iou']:.4f}</td>
                </tr>
                <tr style="background-color: #ffffff;">
                    <td style="padding: 10px;">Avg Proposals/Image</td>
                    <td style="padding: 10px; text-align: right;">{best['avg_proposals_per_image']:.2f}</td>
                </tr>
                <tr style="background-color: #f9f9f9;">
                    <td style="padding: 10px;">Composite Score</td>
                    <td style="padding: 10px; text-align: right;">{best['composite_score']:.4f}</td>
                </tr>
                <tr style="background-color: #ffffff;">
                    <td style="padding: 10px;">Total Images</td>
                    <td style="padding: 10px; text-align: right;">{best['total_images']}</td>
                </tr>
            </table>
        </div>
        """

        display(HTML(html))

    def plot_results(self, figsize=(18, 10)):
        """"""
        if not self.results:
            print(" ")
            return

        viz_path = Path(self.output_dir) / 'optimization_curves.png'

        if viz_path.exists():
            img = Image.open(viz_path)
            plt.figure(figsize=figsize)
            plt.imshow(img)
            plt.axis('off')
            plt.title('Threshold Optimization Results', fontsize=20, fontweight='bold', pad=20)
            plt.tight_layout()
            plt.show()
        else:
            print(f" : {viz_path}")

    def display_recommendations(self):
        """"""
        if not self.results:
            print(" ")
            return

        rec_path = Path(self.output_dir) / 'RECOMMENDATIONS.md'

        if rec_path.exists():
            with open(rec_path, 'r', encoding='utf-8') as f:
                content = f.read()
            display(Markdown(content))
        else:
            print(f" : {rec_path}")

    def apply_best_threshold(self):
        """proposals"""
        if not self.results:
            print(" ")
            return

        best_threshold = self.results['best_threshold']

        print(f" : {best_threshold:.4f}")
        print("=" * 70)

        from src.utils.generate_proposals import run_proposal_generation
        from src.config import STAGE1_CONFIG

        # 
        original_thresh = STAGE1_CONFIG.get('confidence_threshold', 0.1)
        STAGE1_CONFIG['confidence_threshold'] = best_threshold

        try:
            print(" proposals...")
            run_proposal_generation()
            print("\n Proposals !")
            print(f" : proposals/proposals.json")
            print("\n :")
            print("   python main.py train-stage2 --epochs 80")
        except Exception as e:
            print(f"\n : {e}")
            STAGE1_CONFIG['confidence_threshold'] = original_thresh
        finally:
            # 
            pass

    def compare_thresholds(self, show_top_n=5):
        """"""
        if not self.results:
            print(" ")
            return

        import pandas as pd

        all_metrics = self.results['all_metrics']

        # DataFrame
        df = pd.DataFrame([
            {
                'Threshold': m['threshold'],
                'Recall@0.5': m['recall_at_05'],
                'Recall@0.75': m['recall_at_075'],
                'Mean IoU': m['mean_max_iou'],
                'Avg Proposals': m['avg_proposals_per_image'],
                'Composite Score': m['composite_score']
            }
            for m in all_metrics
        ])

        # 
        df_sorted = df.sort_values('Composite Score', ascending=False)

        print(f" Top {show_top_n} :")
        print("=" * 100)

        # 
        display(HTML(df_sorted.head(show_top_n).to_html(index=False, float_format='%.4f')))

        return df_sorted

    def export_summary(self, output_file='threshold_optimization_summary.txt'):
        """"""
        if not self.results:
            print(" ")
            return

        best = self.results['best_metrics']

        summary = f"""
Proposal Threshold Optimization Summary
========================================

Best Threshold: {best['threshold']:.4f}

Performance Metrics:
-------------------
- Recall @ IoU=0.5:  {best['recall_at_05']:.4f} ({best['recall_at_05']*100:.2f}%)
- Recall @ IoU=0.75: {best['recall_at_075']:.4f} ({best['recall_at_075']*100:.2f}%)
- Mean Max IoU:      {best['mean_max_iou']:.4f}
- Avg Proposals/Img: {best['avg_proposals_per_image']:.2f}
- Composite Score:   {best['composite_score']:.4f}

Dataset Info:
------------
- Total Images:    {best['total_images']}
- Total GT Boxes:  {best['total_gt_boxes']}
- Total Proposals: {best['total_proposals']}

Next Steps:
----------
1. Update config.py:
   STAGE1_CONFIG['confidence_threshold'] = {best['threshold']:.4f}

2. Regenerate proposals:
   python main.py gen-proposals --conf-thresh {best['threshold']:.4f}

3. Train Stage 2:
   python main.py train-stage2 --epochs 80

4. Evaluate:
   python main.py evaluate --use-ema
"""

        output_path = Path(self.output_dir) / output_file
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(summary)

        print(f" : {output_path}")
        print("\n" + summary)


# 

def quick_optimize(
    threshold_range=(0.001, 0.1),
    num_points=15,
    sample_size=None
):
    """
    

    Examples:
        >>> results = quick_optimize()
        >>> results = quick_optimize(sample_size=200)  # 
    """
    optimizer = NotebookThresholdOptimizer()
    results = optimizer.run_optimization(
        threshold_range=threshold_range,
        num_points=num_points,
        sample_size=sample_size
    )

    # 
    print("\n")
    optimizer.display_best_result()
    print("\n")
    optimizer.plot_results()

    return results, optimizer


if __name__ == '__main__':
    print("Notebook Threshold Optimizer Helper")
    print("=" * 50)
    print("\n:")
    print("""
    # 1: 
    results, optimizer = quick_optimize()

    # 2: 
    optimizer = NotebookThresholdOptimizer()
    results = optimizer.run_optimization(
        threshold_range=(0.001, 0.1),
        num_points=15,
        sample_size=200
    )
    optimizer.display_best_result()
    optimizer.plot_results()
    optimizer.display_recommendations()
    optimizer.apply_best_threshold()
    """)
